import torch
from torch import nn
import torchvision.transforms as transforms
from uamcf import CompressionManager, PyTorchAdapter
from uamcf.methods import get_method

# 1. 定义/导入模型架构
class YourModelArchitecture(nn.Module):
    # 这里是您的模型架构定义
    pass

# 2. 加载模型
model = YourModelArchitecture()
model.load_state_dict(torch.load("uamcf/model/session0_max_acc.pth"))
model.eval()

# 3. 设置压缩配置
adapter = PyTorchAdapter(device="cuda" if torch.cuda.is_available() else "cpu")
compression_config = {
    "constraints": {
        "accuracy_threshold": 0.9,  # 最低可接受精度(0-1)
        "target_size": 0.25  # 目标大小比例(相对于原始大小)
    },
    "methods": {
        "quantization": {
            "bits": 8,  # 使用8位量化
            "scheme": "asymmetric",
            "per_channel": True
        }
    }
}

# 4. 创建并运行压缩管理器
manager = CompressionManager()
compressed_model, stats = manager.compress(
    model,
    compression_config["constraints"],
    adapter_name="pytorch"
)

# 5. 保存压缩后的模型
torch.save(compressed_model.state_dict(), "保存路径/compressed_model.pth")
print(f"压缩比例: {stats.get('size_reduction', 'N/A')}x")